# coding: utf-8

from typing import Sequence, Tuple
from io import StringIO
from .Graph import vertex_t
from .Subgraph import Subgraph
from .EdgeTypes import LongEdge, UVEdge

# ----------------------------------------------------------------------------------------------
# --------------------------------------  HIERARCHY  -------------------------------------------
# ----------------------------------------------------------------------------------------------

"""
Represents the partitioning of a graph into a set of layers.
Each vertex is in exatly one layer, which is its rank.
"""


class Hierarchy:
    layers: Sequence[Sequence[vertex_t]] = []
    g: Subgraph

    def __init__(self, g: Subgraph, val: int = -1):
        self.g = g
        g.set_all_nodes_attr("rank", val)

    def graph(self):
        return self.g

    """
     Calculates the span of an edge - the number of layers the edge crosses.
     The span can be negative if the edge goes from higher layer to a lower one.
    """

    def span(self, u: vertex_t, v: vertex_t) -> int:
        return self.get_rank(v) - self.get_rank(u)

    def get_rank(self, u: vertex_t) -> int:
        return self.g.get_node_attr(u, "rank")

    def set_rank(self, u: vertex_t, rank: int):
        self.g.set_node_attr(u, "rank", rank)

    def get_pos(self, u: vertex_t) -> int:
        return self.g.get_node_attr(u, "pos")

    def set_pos(self, u: vertex_t, pos: int):
        self.g.set_node_attr(u, "pos", pos)

    # def get_weights(self, u: vertex_t) -> int:
    #     return self.g.get_node_attr(u, 'weights')

    # def set_weights(self, u: vertex_t, weights: int):
    #     self.g.set_node_attr(u, 'weights', weights)

    # number of layers
    def size(self) -> int:
        return len(self.layers)

    # get the layer of the given vertex
    def get_layer(self, u: vertex_t) -> Sequence[vertex_t]:
        return self.layers[self.get_rank(u)]

    # recalculate positions according to the current layers
    def update_all_nodes_pos(self):
        for layer in self.layers:
            i = 0
            for u in layer:
                self.g.set_node_attr(u, "pos", i)
                i += 1

    # get the successor/predecessor of u on its layer
    def next(self, u: vertex_t) -> vertex_t:
        return self.layers[self.get_rank(u)][self.get_pos(u) + 1]

    def prev(self, u: vertex_t) -> vertex_t:
        return self.layers[self.get_rank(u)][self.get_pos(u) - 1]

    # true iff u has a successor/predecessor on its layer
    def has_next(self, u: vertex_t) -> bool:
        return self.get_pos(u) < len(self.get_layer(u)) - 1

    def has_prev(self, u: vertex_t) -> bool:
        return self.get_pos(u) > 0

    # Swap two vertices. Updates all necessary attributes.
    def swap_node(self, u: vertex_t, v: vertex_t):
        self.swap_node_attr(u, v, "pos")
        self.swap_node_attr(u, v, "rank")
        # TODO
        self.layers[self.get_rank(u)][self.get_pos(u)] = u
        self.layers[self.get_rank(v)][self.get_pos(v)] = v

    def swap_node_attr(self, u: vertex_t, v: vertex_t, name: str):
        x, y = self.g.get_node_attr(u, name), self.g.get_node_attr(v, name)
        self.g.set_node_attr(u, name, y)
        self.g.set_node_attr(v, name, x)

    def valid_pos(self, layer_idx: int, i: int) -> bool:
        return i >= 0 and i < len(self.layers[layer_idx])

    """
    Splits all edges that span across more than one layer in the provided Hierarchy
    into segments that each span across just one layer.
    ( aka converts the hierachy into proper Hierarchy )
    @return list of edges which had to be split
    """
    # TODO

    def add_dummy_nodes(self) -> Tuple[Sequence[LongEdge], Sequence[vertex_t]]:
        split_edges: Sequence[LongEdge] = []

        # find edges to be split
        for u in self.g.get_all_nodes():
            for v in self.g.out_neighbours(u):
                span = self.span(u, v)
                if span > 1:
                    split_edges.append(LongEdge(UVEdge(u, v), []))

        # split the found edges
        t_nodes: Sequence[vertex_t] = []
        for se in split_edges:
            orig, path = se.orig, se.path
            span = self.span(orig.u, orig.v)
            path.append(orig.u)

            s: vertex_t = orig.u
            # TODO  span 改为了 span -1
            for i in range(span - 1):
                t = self.g.add_dummy_node()
                t_nodes.append(t)
                self.set_rank(t, self.get_rank(s) + 1)
                # self.g.set_node_attr(t, 'rank', self.get_rank(s)+1)
                self.layers[self.get_rank(t)].append(t)
                self.set_pos(t, len(self.layers[self.get_rank(t)]) - 1)
                # self.g.set_node_attr(t, 'pos', len(
                #     self.layers[self.get_rank(t)]) - 1)

                self.g.add_edge(s, t)
                path.append(t)

                s = t

            path.append(orig.v)
            self.g.add_edge(s, orig.v)
            self.g.remove_edge(orig.u, orig.v)

        return split_edges, t_nodes

    """
    * Counts the number of crossings between edges indident on <u> and <v>
    * if <u> would be to the left of <v>.
    * Takes into account both the ingoing and outgoing edges.
    * <u> and <v> need to be on the same layer
    """

    def crossing_number(self, u: vertex_t, v: vertex_t) -> int:
        count: int = 0
        for out_u in self.g.out_neighbours(u):
            for out_v in self.g.out_neighbours(v):
                if self.get_pos(out_v) < self.get_pos(out_u):
                    count += 1
        for in_u in self.g.in_neighbours(u):
            for in_v in self.g.in_neighbours(v):
                if self.get_pos(in_v) < self.get_pos(in_u):
                    count += 1

        return count

    # counts the number of crossings between layers with index 'layer' and 'layer - 1'
    def count_layer_crossings(self, layer: int) -> int:
        upper = self.layers[layer - 1]
        # lower = h.layers[layer]
        count: int = 0
        for i in range(len(upper) - 1):
            for u in self.g.out_neighbours(upper[i]):
                j = self.get_pos(u)
                for s in range(i + 1, len(upper)):
                    for v in self.g.out_neighbours(upper[s]):
                        t = self.get_pos(v)
                        if (s > i and t < j) or (s < i and t > j):
                            count += 1
        return count

    # counts the total number of crossings in the hierarchy
    def count_crossings(self) -> int:
        count: int = 0
        for i in range(1, self.size()):
            count += self.count_layer_crossings(i)
        return count

    def __str__(self):
        out = StringIO()
        out.write("rank: [")
        for u in self.g.get_all_nodes():
            out.write(f"{u}({self.get_rank(u)}), ")
        out.write("]\n")
        for layer in self.layers:
            for u in layer:
                out.write(f"{u} ")
            out.write("\n")
        return out.getvalue()
